import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torch.nn.functional as F
import pickle

class TextCNN(nn.Module):
    def __init__(self, vocab_size, embedding_size, filter_sizes, num_filters, num_classes):
        super(TextCNN, self).__init__()
        self.embed = nn.Embedding(vocab_size, embedding_size)
        self.conv1ds = nn.ModuleList(
            [nn.Sequential(
                nn.Conv1d(embedding_size, num_filters, kernel_size=filter_sizes[index]),
                nn.ReLU()) for index in range(len(filter_sizes))]
        )
        self.fc1 = nn.Linear(len(filter_sizes)*num_filters, num_classes)

    def forward(self, x):
        out = self.embed(x)
        out = torch.transpose(out, 1, 2)
        out = [conv(out) for conv in self.conv1ds]
        out = [F.max_pool1d(c_out, kernel_size=c_out.size(2)).squeeze(2) for c_out in out]
        out = torch.cat(out, 1)
        logit = self.fc1(out)
        return logit

        # calculate acc
    def get_accuracy(self, logits, target, batch_size):
        """ compute accuracy for training round """
        corrects = (torch.max(logits, 1)[1].view(target.size()).data == target.data).sum()
        accuracy = 100.0 * corrects / batch_size
        return accuracy.item()

